package com.example.demo.config;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.example.demo.dao.StateMachineContextRepository;
import com.example.demo.domain.entity.StateMachineContextEntity;
import com.example.demo.domain.enums.OrderStatus;
import com.example.demo.domain.enums.OrderStatusChangeEvent;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.messaging.MessageHeaders;
import org.springframework.statemachine.StateMachineContext;
import org.springframework.statemachine.StateMachinePersist;
import org.springframework.statemachine.kryo.MessageHeadersSerializer;
import org.springframework.statemachine.kryo.StateMachineContextSerializer;
import org.springframework.statemachine.kryo.UUIDSerializer;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.util.Base64;
import java.util.Objects;
import java.util.UUID;

@Component
public class BizOrderStateMachineContextPersist implements StateMachinePersist<OrderStatus, OrderStatusChangeEvent, String> {

    @Autowired
    private StateMachineContextRepository stateMachineContextRepository;

    private String serialize(StateMachineContext<OrderStatus, OrderStatusChangeEvent> context) {
        Kryo kryo = kryoThreadLocal.get();
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        Output output = new Output(out);
        kryo.writeObject(output, context);
        output.close();
        return Base64.getEncoder().encodeToString(out.toByteArray());
    }

    @SuppressWarnings("unchecked")
    private StateMachineContext<OrderStatus, OrderStatusChangeEvent> deserialize(String data) {
        if (StringUtils.isEmpty(data)) {
            return null;
        }
        Kryo kryo = kryoThreadLocal.get();
        ByteArrayInputStream in = new ByteArrayInputStream(Base64.getDecoder().decode(data));
        Input input = new Input(in);
        return kryo.readObject(input, StateMachineContext.class);
    }

    private static final ThreadLocal<Kryo> kryoThreadLocal = ThreadLocal.withInitial(() -> {
        Kryo kryo = new Kryo();
        kryo.addDefaultSerializer(StateMachineContext.class, new StateMachineContextSerializer());
        kryo.addDefaultSerializer(MessageHeaders.class, new MessageHeadersSerializer());
        kryo.addDefaultSerializer(UUID.class, new UUIDSerializer());
        return kryo;
    });

    @Override
    public void write(StateMachineContext<OrderStatus, OrderStatusChangeEvent> context, String contextObj) {
        StateMachineContextEntity queryResult = stateMachineContextRepository.findTopByOrderId(contextObj);
        if (Objects.nonNull(queryResult)) {
            queryResult.setOrderStatus(context.getState().toString());
            queryResult.setContext(serialize(context));
            stateMachineContextRepository.save(queryResult);
        } else {
            StateMachineContextEntity stateMachineContextEntity = new StateMachineContextEntity();
            stateMachineContextEntity.setOrderId(contextObj);
            stateMachineContextEntity.setOrderStatus(context.getState().toString());
            stateMachineContextEntity.setContext(serialize(context));
            stateMachineContextRepository.save(stateMachineContextEntity);
        }
    }

    @Override
    public StateMachineContext<OrderStatus, OrderStatusChangeEvent> read(String contextObj) {
        StateMachineContextEntity queryResult = stateMachineContextRepository.findTopByOrderId(contextObj);
        if (Objects.nonNull(queryResult)) {
            return deserialize(queryResult.getContext());
        }
        return null;
    }
}
